from OpenGL.GL import glDrawPixels, glPixelZoom, GL_RGBA, GL_UNSIGNED_BYTE
from OpenGL.GLUT import glutCreateMenu, glutAttachMenu, glutAddMenuEntry, GLUT_RIGHT_BUTTON
import pycuda.driver as cuda
from numpy import *
from glwindow import GlWindow
from kernel import Kernel
from param import Param
import time
import Image


class GlGui:
    # members: param_map, img, timer, key, subsample, record, rec_prefix, ortho
    
    def __init__(self, kernel, options):
        self.kernel = kernel
        self.progressive = options.progressive
        self.rec_prefix = options.recording_prefix
        self.param_map = {}
        self.img = None
        self.timer = time.clock()
        self.key = 'Z'
        self.keys = []
        self.subsample = None
        self.record = False
        self.ortho = True
        self.param_resolution = int(options.param_resolution)
        
        
    def start(self):
        # GLUT callbacks
        def init_gl():
            pass

        def menu_func(i):
            self.key = self.keys[i]
            print 'key set to', self.key
            return 0

        def key_func(k):
            if k == 8 and self.rec_prefix != None:
                self.record = not self.record
                print "Recording:", record
            # elif k == '0':
            #     self.ortho = not self.ortho
            elif k >= 'a' and k <= 'z':
                self.key = k.upper()
                print 'key set to:', self.key
                
        def render_func(w, h, matrix, wheel_change, change):
            self.render_scene(w, h, matrix, wheel_change, change)

        # init GLUT, create window
        glw = GlWindow(init_gl, render_func, key_func, "Shadie", True)
        
        # create parameter map
        for p in self.kernel.params:
            if p.bounds != None:
                self.param_map[p.key] = p
                
        # add view parameters (zoom)
        self.param_map['Z'] = Param('zoom', 'float', 1.0, (0.3, 3.0), 'Z', self.param_resolution)
        # self.param_map['D'] = Param('projection_distance', 'float', 1.0, (0.1, 10.0), 'D')

        # create right-click menu
        glutCreateMenu(menu_func)
        i = 0
        
        for p in self.param_map.values():
            glutAddMenuEntry("%s (%s)" % (p.name, p.key), i)
            self.keys.append(p.key)
            i += 1
            
        glutAttachMenu(GLUT_RIGHT_BUTTON)

        # start the display loop
        glw.start()
                
            
    def render_image(self, w, h, subsample, frame, pstruct):
        # run CUDA kernel
        t = time.clock()

        zoom = self.param_map['Z'].value
        dist = -1
        # if self.ortho:
        #     dist = -1
        # else:
        #     dist = self.param_map['D'].value

        self.kernel.run(float32(t - self.timer), int32(w), int32(h), cuda.In(frame),
            float32(zoom), float32(dist),
            cuda.In(pstruct), intp(0), cuda.Out(self.img),
            block = (8, 8, 1), grid = (w/8 + 1, h/8 + 1))
        # print w, time.clock() - t

        # copy data to GL window
        glPixelZoom(self.subsample, self.subsample)
        data = str(self.img.data)
        glDrawPixels(w, h, GL_RGBA, GL_UNSIGNED_BYTE, data)
        
        # record frame to disk
        if self.record:
            f = self.rec_prefix + str(int((t - self.timer)*100)) + '.jpg'
            Image.fromstring('RGBA', (w, h), data).save(f, quality=95)

        
    def render_scene(self, w, h, matrix, wheel_change, change):
        if self.progressive:
            if not change:
                if self.subsample == 1:
                    return
                else:
                    self.subsample /= 2
            else:
                self.subsample = 2 # 4
        else:
            # render every frame, even if no change (for time-dependent views)
            self.subsample = 1
            
        if wheel_change == +1:
            self.param_map[self.key].increment()
        elif wheel_change == -1:
            self.param_map[self.key].decrement()

        # need to resize image buffer?
        if self.img == None or self.img.shape != (h, w, 4):
            self.img = zeros((h, w, 4), dtype = uint8, order = "C")
            
        # setup 4x4 object-to-camera space matrix
        m = hstack([matrix, array([[0], [0], [3]])])
        m = vstack([m, r_[0, 0, 0, 1]])
        
        # invert to get the frame (a camera-to-object matrix)
        frame = linalg.inv(m)
        frame = array(frame[:3, :], order='F', dtype=float32)
        
        # render
        self.render_image(w/self.subsample, h/self.subsample, self.subsample,
            frame, self.kernel.create_params_struct())
        

